#include <iostream>
#include <vector>
#include <numeric>
#include <limits>
#include <algorithm>
#include <omp.h>
#include <random>
#include <assert.h>
#include <sys/time.h>

namespace stackstd {
  template <typename T>
  struct vector {
    T data[33];
    int real_size;
    vector(size_t real_size)
      : real_size(real_size)
    {}

    vector(size_t real_size, T initv)
      : real_size(real_size)
    {
      std::for_each(begin(), end(), [initv](T& t) {t=initv;});
    }

    T* begin() {return &data[0];}
    T* end() {return &data[real_size];}
    const T* begin() const {return &data[0];}
    const T* end() const {return &data[real_size];}

    const T& operator[](int i) const {return data[i];}
    T& operator[](int i) {return data[i];}
  };
}

template <typename T>
struct my_vector {

  typedef T value_type;
  typedef T* iterator;
  typedef T const* const_iterator;

  T* data;
  size_t real_size;

  my_vector(size_t real_size)
    : real_size(real_size)
  {
    data = static_cast<T*>(malloc(real_size*sizeof(T)));
  }

  size_t size() const {return real_size;}

  T* begin() {return &data[0];}
  T* end() {return &data[real_size];}
  const T* begin() const {return &data[0];}
  const T* end() const {return &data[real_size];}

  const T& operator[](int i) const {return data[i];}
  T& operator[](int i) {return data[i];}
};

int nthreads;

template <typename vector_type>
void _stampa(std::string const &s, vector_type &v) {
  // lambdas in openmp regions produce compilation error in closing
  // openmp region... compiler bug??

  typedef typename vector_type::value_type T;

  std::cout << s << ": ";
  std::for_each(v.begin(), v.end(), [](T t) {std::cout << t << " ";});
  std::cout << std::endl;
}

#define stampa(a) _stampa( std::string(#a), a);

#define PR(a) #a << " = " << a

double lapse_time;
double lapse_time_seq;

template <typename vector_type>
bool check_sort(vector_type & a, vector_type & b) {
  // true if fails

  typedef typename vector_type::value_type T;

  struct timeval start_tv;
  struct timeval stop_tv;

  gettimeofday(&start_tv, NULL);
  std::sort(a.begin(), a.end());
  gettimeofday(&stop_tv, NULL);

  lapse_time_seq = ((static_cast<double>(stop_tv.tv_sec)+1/1000000.0*static_cast<double>(stop_tv.tv_usec)) - (static_cast<double>(start_tv.tv_sec)+1/1000000.0*static_cast<double>(start_tv.tv_usec))) * 1000.0;

  std::cout << "\nChecking results\n" << std::endl;
  // stampa(a);
  // stampa(b);

  bool res = std::equal(a.begin(), a.end(), b.begin());
  std::cout << std::boolalpha << res << std::endl;
  if (!res) {
  	std::pair<typename vector_type::iterator, typename vector_type::iterator> value;
  	value = std::mismatch(a.begin(), a.end(), b.begin());
  	std::cout << "First element that fails is "
              << *value.first << " " << * value.second << " "
              << "distance " << value.first-a.begin()
              << std::endl;
  }
						 
  return !res;
}

template <typename vector_type>
void sample_sort(vector_type const & a, vector_type & b) {

  typedef typename vector_type::value_type T;

  std::random_device rd;
  std::mt19937 gen(rd());
  std::uniform_int_distribution<> dis(0, a.size()-1);

  std::vector<T> samples(nthreads+1);
  samples[0] = std::numeric_limits<T>::min();
  samples[nthreads] = std::numeric_limits<T>::max();

  std::vector<int> balance(nthreads);

#ifndef NDEBUG
  std::cout << std::numeric_limits<T>::min() << " "
            << std::numeric_limits<T>::max() << std::endl;
#endif
  for (int i = 1; i < nthreads; ++i)
    samples[i] = a.size()/nthreads*i;//a[dis(gen)];

  std::sort(samples.begin(), samples.end());

#ifndef NDEBUG
  std::cout << "Samples: ";
  stampa(samples);
  std::cout << std::endl;
#endif

  std::vector<std::vector<int>*> pcounts(nthreads);
  std::vector<std::vector<int>*> pprefix(nthreads);

  std::vector<int> binbegin(nthreads);
  std::vector<int> binend(nthreads);

#pragma omp parallel
  {
    int th = omp_get_thread_num();
    const int begin = th * (a.size()/nthreads);
    const int end = (th+1) * (a.size()/nthreads);

#ifndef NDEBUG
#pragma omp critical
    {
      std::cout << "begin " << begin << " "
                << "end "   << end << " "
                << "((end-begin)*nthreads)==a.size()? " 
                << std::boolalpha << (((end-begin)*nthreads)==a.size())
                << std::endl;
    }
#endif
    assert((((end-begin)*nthreads)==a.size()));

    std::vector<int> counts(nthreads);
    std::vector<int> prefix(nthreads);
    std::vector<int> current(nthreads,0);
    pcounts[th] = &counts;
    pprefix[th] = &prefix;

    for (int i = begin; i < end; ++i) {
      int t = 1;
      while (t < nthreads && a[i] > samples[t]) {
        // #pragma omp critical
        // 		{
        // 		  std::cout << th << " -> * " 
        // 					<< t << " === "
        // 					<< a[i] << "," << samples[t+1]
        // 					<< std::endl;
        // 		}
        ++t;
      }
      // #pragma omp critical
      // 		{
      // 		  std::cout << th << " -> **** " << t << std::endl;
      // 		}
      if (t <= nthreads)
        counts[t-1]++;
    }
#ifndef NDEBUG
#pragma omp critical
    {
      std::cout << th << "-> " ;
      //                << begin << " " << end << std::endl;
      stampa(counts);
    }
#endif
#pragma omp barrier
#pragma omp single
    {
      // prefix sums done by one thread
      int start = 0;
      for (int j=0; j<nthreads; ++j) {
        (*(pprefix[0]))[j] = start;
        for (int i=1; i<nthreads; ++i) {
          (*(pprefix[i]))[j] = (*(pprefix[i-1]))[j] + (*(pcounts[i-1]))[j];
        }
        start = (*(pprefix[nthreads-1]))[j] + (*(pcounts[nthreads-1]))[j];
      }
    } // end single

    for (int i=0; i<nthreads; ++i) {
      binbegin[i] = (*(pprefix[0]))[i];
      binend[i] = (i==nthreads-1)?a.size():(*(pprefix[0]))[i+1];
    }
    //#pragma omp barrier ci deve gia` essere

#ifndef NDEBUG
#pragma omp critical
    {
      std::cout << th << "-> " ;
      stampa(prefix);
    }
#endif
    for (int i = begin; i < end; ++i) {
      int t = 1;
      while (t < nthreads && a[i] > samples[t]) {
        ++t;
      }
      //t-1 e` il bin dove mettere i dati
      b[prefix[t-1]+current[t-1]] = a[i];
      ++current[t-1];
    }

    // #pragma omp critical
    //     {
    //       std::cout << th << "-> " ;
    //       stampa(b);
    //     }
#pragma omp barrier
    // int binbegin = prefix[0], binend;
    // if (th==nthreads-1)
    //   binend=a.size();
    // else
    //   binend=prefix[th+1];
#ifndef NDEBUG
#pragma omp critical
    {
      std::cout << th << "-> "
                << PR(binbegin[th]) << " "
                << PR(binend[th]) << " "
                << std::endl;
    }
#endif
    balance[th] = binend[th] - binbegin[th];

    std::sort(&(b[binbegin[th]]), &(b[binend[th]]));
#pragma omp barrier
    // #pragma omp critical
    //     {
    //       std::cout << th << "-> " ;
    //       stampa(b);
    //     }
  } // omp parallel

  std::cout << "Balance " 
            << (static_cast<float>(*std::max_element(balance.begin(), balance.end())) /
                static_cast<float>(*std::min_element(balance.begin(), balance.end())))
            << std::endl;

}

#define TYPE int

int main(int argc, char** argv) {
  struct timeval start_tv;
  struct timeval stop_tv;

  int N=atoi(argv[1]);
  //std::numeric_limits<TYPE>::min(),
  //std::numeric_limits<TYPE>::max());

  nthreads = omp_get_max_threads();
  std::cout << "Threads " << nthreads << std::endl;

  gettimeofday(&start_tv, NULL);
#ifdef FIRST_TOUCH 
  my_vector<TYPE> data(N);
  my_vector<TYPE> out(N);
#pragma omp parallel
  {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0,N);
    int th = omp_get_thread_num();
    int local_size = N/nthreads;
    std::for_each(&data[local_size*th], &data[local_size*(th+1)],  [&dis, &gen](TYPE& x) {x = dis(gen);});
    std::for_each(&out[local_size*th], &out[local_size*(th+1)],  [](TYPE& x) {x = 0;});
  }
#else
  std::random_device rd;
  std::mt19937 gen(rd());
  std::uniform_int_distribution<> dis(0,N);

  std::vector<TYPE> data(N);
  std::vector<TYPE> out(N);

  std::for_each(data.begin(), data.end(), [&dis, &gen](TYPE& x) {x = dis(gen);});
#endif
  gettimeofday(&stop_tv, NULL);
  double init_time = ((static_cast<double>(stop_tv.tv_sec)+1/1000000.0*static_cast<double>(stop_tv.tv_usec)) - (static_cast<double>(start_tv.tv_sec)+1/1000000.0*static_cast<double>(start_tv.tv_usec))) * 1000.0;

#ifndef NDEBUG
  stampa(data);
#endif

  gettimeofday(&start_tv, NULL);
  sample_sort(data, out);
  gettimeofday(&stop_tv, NULL);

  lapse_time = ((static_cast<double>(stop_tv.tv_sec)+1/1000000.0*static_cast<double>(stop_tv.tv_usec)) - (static_cast<double>(start_tv.tv_sec)+1/1000000.0*static_cast<double>(start_tv.tv_usec))) * 1000.0;

  bool res =  check_sort(data, out);

  std::cout << "Time: " << lapse_time << "ms" 
            << " Speed-up: " << lapse_time_seq/lapse_time
            << " (seq time: " << lapse_time_seq << ")"
            << " init: " << init_time
            << std::endl;

  std::cout << "grep " 
            << lapse_time << " "
            << lapse_time_seq/lapse_time << " "
            << lapse_time_seq << " "
            << lapse_time * nthreads << " "
            << lapse_time * nthreads - lapse_time_seq
            << std::endl;
    
  return res;
}
